from flask import Flask, request, jsonify, Response
from flask_cors import CORS
import json
import os
from werkzeug.utils import secure_filename
from models.response import ResponseMessage
from models.enums import StatusCode
from ConfigManager import ConfigManager
from models.enums import EventMessageType
import json
from utils.RedisQueue import RedisQueue
from utils.WaitUtils import WaitUtils
from utils.MessageEnrichUtils import MessageEnrichUtils
from models.enums import EventMessageType
from utils.RagTaskManger import RagTaskManger
from Logger import Logger


logger = Logger.get_logger(__name__)

app = Flask(__name__)
max_file_upload_size = ConfigManager().get("max_file_upload_size")
# 配置文件上传参数
if max_file_upload_size:
    app.config["MAX_CONTENT_LENGTH"] = int(max_file_upload_size) * 1024 * 1024
else:
    app.config["MAX_CONTENT_LENGTH"] = 20 * 1024 * 1024

# 启用 CORS
CORS(
    app,
    resources={
        r"/api/*": {
            "origins": ["*"],  # 允许的域名
            "methods": ["GET", "POST", "OPTIONS"],  # 允许的方法
            "allow_headers": ["*"],  # 允许的请求头
            "supports_credentials": True
        },
        r"/*": {
            "origins": "*",  # 允许的域名
            "methods": ["GET", "POST", "OPTIONS"],  # 允许的方法
            "allow_headers": ["access-control-allow-origin","content-type"],  # 允许的请求头
        },
    },
)



def decode_request_data(raw_data):
    # 尝试不同的编码方式
    encodings = ["utf-8", "gbk", "gb2312", "iso-8859-1"]
    data = None

    for encoding in encodings:
        try:
            if isinstance(raw_data, bytes):
                decoded_data = raw_data.decode(encoding)
            else:
                decoded_data = raw_data
            data = json.loads(decoded_data)
            break
        except Exception as e:
            logger.error(f"Failed to decode with {encoding}: {str(e)}")
            continue

    return data


def generate_watch(uniq_id: str):
    """Generate SSE events for client watching"""
    while True:
        try:
            message = RedisQueue().dequeue(EventMessageType.FRONT.value, uniq_id)
            if not message:
                WaitUtils.wait(0.5)
                continue
            if not message.get("response"):
                WaitUtils.wait(0.5)
                continue
            data = json.dumps(message)
            yield f"data: {data}\n\n"
        except Exception as e:
            logger.error(f"Error while consuming message: {str(e)}")
            yield f"data: {json.dumps({'error': str(e)})}\n\n"


@app.route("/api/watch", methods=["POST"])
def watch():
    raw_data = request.get_data()
    data = decode_request_data(raw_data)
    if data is None:
        return jsonify({"error": "Unable to decode request data"}), 400

    if not data or "uniq_id" not in data:
        return jsonify({"error": "Missing messages parameter"}), 400
    uniq_id = data.get("uniq_id")
    logger.info(f"{uniq_id} connected to watch")
    response = Response(generate_watch(uniq_id), mimetype="text/event-stream")
    response.headers["Cache-Control"] = "no-cache"
    response.headers["Connection"] = "keep-alive"
    return response


@app.route("/api/watch/<uniq_id>", methods=["GET"])
def watch_get(uniq_id):
    logger.info(f"{uniq_id} connected to watch")
    return Response(generate_watch(uniq_id), mimetype="text/event-stream")


@app.route("/api/upload", methods=["POST"])
def upload_file():
    """文件上传接口"""
    try:
        # 检查是否有文件
        if "file" not in request.files:
            return ResponseMessage.error(
                "未找到上传文件", code=StatusCode.BAD_REQUEST
            ).model_dump_json()

        file = request.files["file"]
        if file.filename == "":
            return ResponseMessage.error(
                "未选择文件", code=StatusCode.BAD_REQUEST
            ).model_dump_json()

        # 获取文件保存目录
        upload_dir = ConfigManager().get("temp_file_dir")
        if not upload_dir:
            return ResponseMessage.error(
                "服务端未配置文件保存目录", code=StatusCode.SERVER_ERROR
            ).model_dump_json()

        # 确保目录存在
        os.makedirs(upload_dir, exist_ok=True)

        # 安全处理文件名并保存文件
        # secure_filename(file.filename)
        filename = file.filename
        file_path = os.path.join(upload_dir, filename)
        file.save(file_path)

        # 发布文件上传成功事件
        user_message = {"file_path": file_path, "file_name": file.filename}
        _, user_message = MessageEnrichUtils.enrich_uniq_message(user_message)
        user_message = MessageEnrichUtils.enrich_handler(user_message,EventMessageType.FILE.value)
        RagTaskManger.start_task(user_message)
        # 打印文件路径
        logger.info(f"文件已保存: {file_path}")
        return ResponseMessage.success(
            data=None, msg=f"{filename}文件上传成功"
        ).model_dump_json()

    except Exception as e:
        logger.error(f"文件上传失败: {str(e)}")
        return ResponseMessage.error(
            msg=f"文件上传失败: {str(e)}", code=StatusCode.SERVER_ERROR
        ).model_dump_json()


@app.route("/api/rag", methods=["POST"])
def generate():
    # 添加调试输出
    logger.info(f"Received request data: {request.get_data()}")
    try:
        # 获取原始数据并解码
        raw_data = request.get_data()
        data = decode_request_data(raw_data)

        if data is None:
            return jsonify({"error": "Unable to decode request data"}), 400

        if not data or "messages" not in data:
            return jsonify({"error": "Missing messages parameter"}), 400

        messages = data.get("messages", [])
        if not messages or not isinstance(messages, list):
            return jsonify({"error": "Invalid messages format"}), 400

        user_message = next(
            (msg for msg in messages if msg.get("role") == "user"), None
        )
        if not user_message or "content" not in user_message:
            return jsonify({"error": "Missing user message content"}), 400
        uniq_id, user_message = MessageEnrichUtils.enrich_uniq_message(user_message)
        user_message = MessageEnrichUtils.enrich_handler(user_message,EventMessageType.RAG.value)
        RagTaskManger.start_task(user_message)
        return jsonify({"success": "task received", "uniq_id": uniq_id}), 200
    except Exception as e:
        logger.error(f"Error: {str(e)}")  # 添加错误输出
        return jsonify({"error": str(e)}), 400


if __name__ == "__main__":
    app.run(host="127.0.0.1", port=8088, debug=True)
